************************************************************
* Treatment effect heterogeneity by predicted dosage
************************************************************
clear all
set more off

// Install necessary packages
net install dm79.pkg, from(https://www.stata.com/stb/stb56/)

// Read in command line arguments
args covs roster estrat_dose output output_dir

* Load local parameter called covariates which is created using ***
include "`covs'"

* Load cross-sectional READI analysis sample data
insheet using "`roster'", comma clear

* Load modified endogeneous stratification function (estrat_dose)
do "`estrat_dose'"

* Prepare data
destring t_payroll_total_hours_post_20, force replace
replace t_payroll_total_hours_post_20=. if treatment==0

gen pay_dose=t_payroll_total_hours_post_20
replace pay_dose=0 if treatment==1 & t_headcount_post_20==0

rename i_three_components_post_20 index
rename a_shooting_or_homicide_post_20 ashoot
rename a_p1_non_sh_post_20 aother
rename v_shooting_or_homicide_post_20 vshoot

* Use continuous versions of baseline measures + strata fixed effects

foreach k in index ashoot aother vshoot {
	di "outcome = `k'"

	set seed 9998405
	sort cluster

	estrat_dose `k' treatment, dosage(pay_dose) pred(`baselines' `strataFE') cov(`baselines' `strataFE') groups(2) reps(500) boot(500) savegroup

	matrix `k' = e(LOO_C)',e(LOO_SE)',e(RSS_C)',e(RSS_SE)'
	matrix rownames `k' = group1 group2
	matrix colnames `k' = LOO_B LOO_SE RSS_B RSS_SE
}

* Save output with LOO group identifiers
save "`output'", replace

* Save outcome-specific estimates
foreach k in index ashoot aother vshoot {
	preserve
	clear
	svmat2 `k', names(LOO_B LOO_SE RSS_B RSS_SE)
	gen group=_n
	gen outcome="`k'"
	order outcome group
	save "`output_dir'dosage_heterogeneity_`k'.dta", replace
	restore
}

